from typing import Tuple
import torch
from loadDataset import MakeDataSet
import numpy as np


class BCICompDatasetForMulti(torch.utils.data.Dataset):
    def __init__(self, channelNum: Tuple[int, int, int], subjectARootPath: str, subjectBRootPath: str,
                 convertToF: bool = True):
        super().__init__()
        self.A_tarDatas, self.A_noTarDatas = MakeDataSet.LoadTarAndNoTarData(subjectARootPath, convertToF)
        self.B_tarDatas, self.B_noTarDatas = MakeDataSet.LoadTarAndNoTarData(subjectBRootPath, convertToF)
        A_tarLen = int(np.size(self.A_tarDatas) / np.size(self.A_tarDatas, -1))
        A_noTarLen = int(np.size(self.A_noTarDatas) / np.size(self.A_noTarDatas, -1))
        b_tarLen = int(np.size(self.B_tarDatas) / np.size(self.B_tarDatas, -1))
        b_noTarLen = int(np.size(self.B_noTarDatas) / np.size(self.B_noTarDatas, -1))
        self.dataLen = A_tarLen + A_noTarLen + b_tarLen + b_noTarLen
        self.datas = np.zeros((self.dataLen, 6, 240))
        self.labels = np.zeros(self.dataLen, dtype=np.int16)
        count = 0
        for i in range(np.size(self.A_tarDatas, 0)):
            A_electrodes = np.array(
                [self.A_tarDatas[i, channelNum[0], :, :, :], self.A_tarDatas[i, channelNum[1], :, :, :],
                 self.A_tarDatas[i, channelNum[2], :, :, :]])
            B_electrodes = np.array(
                [self.B_tarDatas[i, channelNum[0], :, :, :], self.B_tarDatas[i, channelNum[1], :, :, :],
                 self.B_tarDatas[i, channelNum[2], :, :, :]])
            for j in range(np.size(A_electrodes, 1)):
                for m in range(np.size(A_electrodes, 2)):
                    if np.all(A_electrodes[0, j, m, :] == 0):
                        continue
                    self.datas[count, 0:3, :] = A_electrodes[:, j, m, :]
                    self.datas[count, 3:6, :] = B_electrodes[:, j, m, :]
                    self.labels[count] = 1
                    count += 1
        for i in range(np.size(self.A_noTarDatas, 0)):
            A_electrodes = np.array(
                [self.A_noTarDatas[i, channelNum[0], :, :, :], self.A_noTarDatas[i, channelNum[1], :, :, :],
                 self.A_noTarDatas[i, channelNum[2], :, :, :]])
            B_electrodes = np.array(
                [self.B_noTarDatas[i, channelNum[0], :, :, :], self.B_noTarDatas[i, channelNum[1], :, :, :],
                 self.B_noTarDatas[i, channelNum[2], :, :, :]])
            for j in range(np.size(A_electrodes, 1)):
                for m in range(np.size(A_electrodes, 2)):
                    if np.all(A_electrodes[0, j, m, :] == 0):
                        continue
                    self.datas[count, 0:3, :] = A_electrodes[:, j, m, :]
                    self.datas[count, 3:6, :] = B_electrodes[:, j, m, :]
                    self.labels[count] = 0
                    count += 1

    def __len__(self) -> int:
        return self.dataLen

    def __getitem__(self, idx) -> Tuple[np.array, np.int16]:
        if idx > self.dataLen:
            print(idx)
            return self.datas[self.dataLen - 1, :, :], self.labels[self.dataLen - 1]
        return self.datas[idx, :, :], self.labels[idx]
